# import matplotlib.pyplot as plt
import numpy as np
# from utils.data_to_color import ImColor_tp
from scipy.interpolate import griddata
import glob
from utils.config import station_grid_path, label_path
import os
from multiprocessing import Pool, cpu_count
import json

# 组合反射率区域: latmin, latmax, lonmin, lonmax = 19, 28, 102, 114
# apcp全区域:  18.21, 27.8, 102.2, 113.79
# 实际选取区域: 19.01, 27.8, 102.2, 113.79
# latmax, latmin, lonmin, lonmax = 28, 19.01, 102, 113.99 #[20:, 20:-20], [120:-200, 200:-160]
latmin, latmax, lonmin, lonmax = 21.01, 26.6, 104.2, 112.19
xdim, ydim, reso = 560, 800, 0.01


def interp(file):
    print(file)
    _, fname = os.path.split(file)

    station_data = np.load(file)[140:-200, 220:-180]
    # points = np.where(station_data > 0)
    values = station_data[idxs]

    grid_lat = np.arange(xdim)
    grid_lon = np.arange(ydim)
    lon_grid, lat_grid = np.meshgrid(grid_lon, grid_lat)

    lonl = lon_grid[idxs]
    latl = lat_grid[idxs]
    _points = (lonl, latl)

    # nearest = griddata(_points, values, (lon_grid, lat_grid), method='nearest', rescale=True)
    # linear = griddata(_points, values, (lon_grid, lat_grid), method='linear', rescale=True)
    cubic = griddata(_points, values, (lon_grid, lat_grid), method='cubic', fill_value=0, rescale=True)

    pxs, pys = [], []
    st = 15
    for x in range(xdim):
        for y in range(ydim):
            x0 = max(0, x - st)
            x1 = min(xdim, x + st)
            y0 = max(0, y - st)
            y1 = min(ydim, y + st)

            era = station_data[x0:x1, y0:y1] > 0

            if era.sum() < 1:
                pxs.append(x)
                pys.append(y)
    cubic[([pxs], [pys])] = 0
    cubic[cubic < 0] = 0
    cubic = np.nan_to_num(cubic)

    sdir = f"{label_path}/{fname[:8]}"
    if not os.path.exists(sdir):
        os.mkdir(sdir)
    np.savetxt(f"{sdir}/{fname[:12]}.000", cubic, fmt='%.02f')


if __name__ == '__main__':
    with open('/home/gym/projects/my_test/data/idxs.json') as f:
        idxs = tuple(np.array(json.load(f)).T)
    files = glob.glob(f"{station_grid_path}/20211*.npy")
    for f in files:
        interp(f)
        exit()

    # pool = Pool(processes=20)
    # pool.map(interp, files)
    # pool.close()
    # pool.join()
